I will suggest using a customizable pre-hook for creating the dummy tensors. Here is the PR for adding this feature:
#6924
We can replace this change with the pre-hook later.

The prehook in PR#6924 is useful and we should use it once that is merged. I see a need to initialize data differently under different precisions due to the dynamic range difference.

Adding @nekorobov for vis. When #6924 is merged, we may want to adjust how dummy tensors are initialized during the autotune under different precisions. In order to have ideal mix of 0s and 1s.

#6924 has been merged. Because using random data instead of all-zero tensors can be considered an improvement for the perf stabilization in general, I think we can also keep the code change here. Thanks a lot for the suggestion~

For this part, I want to confirm:

Currently, for this series of moe runners, tile_tokens_dim is created according to the value of num_tokens through the calculate_tile_tokens_dim method. Thus, when tuning num_tokens, we should determine the tile_tokens_dim after a specific num_tokens is given.

Now we tune each num_tokens with all the possible tile_tokens_dim. But this only stands when the two values are independent of each other. Has it been decoupled already?

But I still see calculate_tile_tokens_dim is used in the fake registration part and some other places, which means tile_tokens_dim is still a function of num_tokens. Did I miss anything here?

But this only stands when the two values are independent of each other. Has it been decoupled already?

But I still see calculate_tile_tokens_dim is used in the fake registration part and some other places, which means tile_tokens_dim is still a function of num_tokens. Did I miss anything here?

The ideal tile_tokens_dim correlates to num_tokens and outside factors such as expert distribution. We don't have good heurstics yet. Ideally, we may predict a range of tile_tokens_dim according to num_tokens to reduce tuning space.

Is my understanding correct that register_fake() must return correct output shapes for torch.compile to work? If so, I see an possible issue--the output shape max_num_padded_tokens depends on tile_tokens_dim in current implementation, but we cannot determine tile_tokens_dim by heuristics because it's tuned by the autotuner. Is there any side effect to nail down tile_tokens_dim=128 in register_fake()?

Is my understanding correct that register_fake() must return correct output shapes for torch.compile to work?

I believe so. The fake part will do the shape inference, so it is expected to have the exact same shape as the custom op implementation.

A possible solution might be always using the max tile value (like 128) to pad the tensor, which will be easier for the fake part. Not sure whether it will be suboptimal or not.

And I wonder how the following op uses this padded tensor, as its shape has changed. Is there any extra information delivered?

When autotuner is enabled, the true tile_tokens_dim can be obtained from the return tuple (tile_tokens_dim, tactic) of choose_one() (which I just added today in de91deb). Or, when autotuner is disabled, estimate the tileN using the same heuristic. I haven't implemented both. Does it only affect torch.compile? If so, I'd like to have a follow-up PR for that.

I think only torch.compile will be affected.

If I understand correctly, creating multiple runners is for tuning the tileN, and tile_tokens_dim can be independent of the num_tokens (this is important), then I suggest using tactics to represent them instead of different runners. Because we prefer using different runners to represent a different backend or implementation, and different tactics in a single runner to represent a different kernel config.

A proper modification can be:

• 'Encode` it into different tactics and prepare all the CPP runner instances in a single Python runner.
• In 'get_valid_tactics', create a production between tactic list and possible tileN values. Return a tuple (tile_tokens_dim, tactic) instead of a standalone tactic value to indicate the corresponding runner used for that tile value.
• In forward, invoke the correct runner according to the tile_tokens_dim passed within the tactic tuple.

Sounds reasonable. Let me revise it

Added in de91deb

#6924 has been merged. Because using random data instead of all-zero tensors can be considered an improvement for the perf stabilization in general, I think we can also keep the code change here. Thanks a lot for the suggestion~

Actually, I think we can do it directly in the Python code because it might be simpler. But CPP here is also fine.

@nekorobov what's your opinion on this?

I do not have strong opinion. Whatever is easier and faster in your opinion :) Since CPP is already implemented I am fine having it this way

I don't have strong preference as well. Let's keep it as-is (cpp) so python side is cleaner.

Can tactics be shared by these runners with different values of tile_tokens_dim? Because we calculate this value only according to the input number of tokens. During the warm-up phase, it is set to be the maximum number of tokens. Then the tuning process will always rely on the runner with that value. Is this expected? Some potential issues might be:

• Using a runner with the large tile to tune a small input problem size, which might be suboptimal.
• Using tactics across different tile values. Some runners might not support others' tactics. Not sure if the test failures mentioned in the comment note is caused by this.

I think autotuner tunes the "dynamic dimension" of the input tensors from 1, 2, 4, ..., up to some large max_batch size during the warm-up phase. So all shapes should be covered by autotuner.

Tactics are not interchangeable between runners. So for this MoE precision, we don't tune tileN due to an issue I met earlier. As result, the autotuner has no visibility into who is the best runner. It only sees the best tactic. calculate_tile_tokens_dim() fills the gap and chooses the right runner because the result is deterministic.